################################################################################
class rap_states_jj(rap_states):
  """ Store set of JJ states. """
  def __init__(self,filename=None,**args):
    self.clear()
    if filename: self.loadfile(filename,**args)

# ------------------------------------------------------------------------------

  def _reprocess(self):
    """ Redefine senority-like tags. """
    if len(self) == 0: return
    indices=zip(*sorted(zip(self.v_energy.values(),self.v_energy.keys())))[1]
    allowedterms={}

    have={}
    tmp={}
    for idx in indices:
      (conf,j2,sen,par,en)=self.level(idx)

      key=(conf,j2)
      if key not in have: have[key]=0
      have[key]+=1
      self.v_seniority[idx]=have[key]

# ------------------------------------------------------------------------------

  def _map_jj(self,other,conf=False,alpha=False,common=True,minimize=False,\
              **args):
    """ 
    Set level map to rap_states_jj instance. 

    Setting alpha to True (default) means that the level labels can be trusted
    (conf and common are assumed True).  Otherwise, if conf is True, then 
    configuration is assumed to be a good quantum number and mapping is done
    over (conf,J) pairs; mapping is done in energy order unless minimize is set
    to True when energy differences are minimized (only really necessary
    when number of levels per symmetry is different).  Finally, if conf and
    alpha are False and common is True, then mapping is done over states 
    belonging to the common set of configurations and according to (J,p)
    symmetries.  If conf, alpha, and common are all False, then mapping is
    done over all states according to (J,p).  Setting minimize=True in this 
    case will attempt to minimize energy differences between states.
    """
    if alpha is True:         # seniority numbers are trusted
      tmap={}
      for i in other:
        k=None
        for j in self:
          if self[j] == other[i]:
            k=j
            break
        tmap[i]=k
      self.v_lmap=tmap
      self.v_lmap_r=dict([(v,k) for k,v in self.v_lmap.items()])
      return

    fromidxs=self.indices()
    toidxs=other.indices()
    if common is True:
      fromconfs=set(self.configurations())
      toconfs=set(other.configurations())
      cconfs=list(fromconfs.intersection(toconfs))
      fromidxs=self.cfilter(cconfs)
      toidxs=other.cfilter(cconfs)

    if conf is True:
      fromlist={}
      for i in fromidxs:
        key=(self.configuration(i),self.j2(i))
        if key not in fromlist: fromlist[key]=[]
        fromlist[key].append(i)
      tolist={}
      for j in toidxs:
        key=(other.configuration(j),other.j2(j))
        if key not in fromlist: continue
        if key not in tolist: tolist[key]=[]
        tolist[key].append(j)
    else:
      fromlist={}
      for i in fromidxs:
        key=(self.j2(i),self.parity(i))
        if key not in fromlist: fromlist[key]=[]
        fromlist[key].append(i)
      tolist={}
      for j in toidxs:
        key=(other.j2(j),other.parity(j))
        if key not in fromlist: continue
        if key not in tolist: tolist[key]=[]
        tolist[key].append(j)

    if minimize is True:
      self.v_lmap=self._map_by_energy(other,fromlist,tolist)
    else:
      self.v_lmap=self._map_straight(other,fromlist,tolist)
    for j in other:
      if j not in self.v_lmap: self.v_lmap[j]=None
    self.v_lmap_r=dict([(v,k) for k,v in self.v_lmap.items() if v is not None])
    for i in self:
      if i not in self.v_lmap_r: self.v_lmap_r[i]=None

# ------------------------------------------------------------------------------

  def _map_lsj(self,other,conflist=None,**args):
    """ 
    Set level map to rap_states_lsj instance. 

    Mapping is performed only over the levels sharing a set of common
    configurations.  The mapping is based on J, parity, and energy order.
    """
    if conflist is not None:
      if isinstance(conflist[0],rap_config):
        oconfs=set(conflist[:])
        mejconfs=set(self.configurations())
      elif isinstance(conflist[0],rap_jconfig):
        oconfs=set(other.configurations())
        mejconfs=set(conflist[:])
      else:
        msg='confset must contains instances of rap_config or rap_jconfig'
        raise ErrorRAP(msg)
    else:
      oconfs=set(other.configurations())
      mejconfs=set(self.configurations())

    meconfs=[]
    for jconf in mejconfs:
      conf=jconf.to_config()
      if conf not in meconfs: meconfs.append(conf)
    meconfs=set(meconfs)
    shared=list(oconfs & meconfs)
    jshared=[]
    for conf in shared: jshared.extend(conf.to_jconfig())
    oidxs=other.cfilter(shared)
    meidxs=self.cfilter(jshared)

    (j2min,j2max)=other.j2range(levlist=oidxs)
    pars=[0,1]
    fromlist={}
    tolist={}
    for p in pars:
      for j2 in range(j2min,j2max+1,2):
        sym=(j2,p)
        toidxs=sorted(other.symfilter(sym,subset=oidxs))
        tmeidxs=sorted(self.symfilter(sym,subset=meidxs))
        if len(toidxs) != len(tmeidxs):        # sanity check
          msg='level sets differ in number of levels for symmetry: '+\
              '2J,p=%i,%i'%(j2,p)
          raise ErrorRAP(msg)
        fromlist[sym]=tmeidxs
        tolist[sym]=toidxs

    self.v_lmap=self._map_by_energy(other,fromlist,tolist)
    for j in other:
      if j not in self.v_lmap: self.v_lmap[j]=None
    self.v_lmap_r=dict([(v,k) for k,v in self.v_lmap.items() if v is not None])
    for i in self:
      if i not in self.v_lmap_r: self.v_lmap_r[i]=None

# ------------------------------------------------------------------------------

  def _map_jp(self,other,conflist=None,**args):
    """ 
    Set level map to rap_states_jp instance. 

    Mapping is performed only over all levels and based on J, parity, and 
    energy order.
    """
    idxs=self.indices()
    if conflist is not None:
      if isinstance(conflist[0],rap_config):
        tconflist=[]
        for conf in conflist: tconflist.extend(conf.to_jconfig())
        conflist=tconflist
      idxs=self.cfilter(conflist)
    (j2min,j2max)=self.j2range(levlist=idxs)
    pars=[0,1]
    fromlist={}
    tolist={}
    for p in pars:
      for j2 in range(j2min,j2max+1,2):
        sym=(j2,p)
        fromlist[sym]=self.symfilter(sym)
        tolist[sym]=other.symfilter(sym)

    self.v_lmap=self._map_by_energy(other,fromlist,tolist)
    for j in other:
      if j not in self.v_lmap: self.v_lmap[j]=None
    self.v_lmap_r=dict([(v,k) for k,v in self.v_lmap.items() if v is not None])
    for i in self:
      if i not in self.v_lmap_r: self.v_lmap_r[i]=None

# ------------------------------------------------------------------------------

  def clear(self):
    """ Erase all data. """
    rap_states.clear(self)
    self.v_cup='JJ'
    self.v_jconf={}
    self.v_j2={}
    self.v_seniority={}

# ------------------------------------------------------------------------------

  def clear_levels(self):
    """ Just erase level data. """
    rap_states.clear_levels(self)
    self.v_cup='JJ'
    self.v_jconf={}
    self.v_j2={}
    self.v_seniority={}

# ------------------------------------------------------------------------------

  def loadfile(self,filename,**args):
    """ Load in level information from level data file. """
    ft=rlf_filetype(filename)
    if ft == 'fac_lev':
      self.load_fac(filename,**args)
    else:
      msg='do not know how to read file: %s'%filename
      raise ErrorRAP(msg)

# ------------------------------------------------------------------------------

  def reindex(self,remap=None):
    """ 
    Reassign level indices.  Indices are remapped according to dictionary,
    remap, old (key) -> new (value).  If remap is None, then reindex according
    to energy order.
    """
    imap=rap_states.reindex(self,remap)
    self.v_jconf=dict([(imap[i],val) for (i,val) in self.v_jconf.items()])
    self.v_j2=dict([(imap[i],val) for (i,val) in self.v_j2.items()])
    self.v_seniority=dict([(imap[i],val) for (i,val) in \
                          self.v_seniority.items()])

# ------------------------------------------------------------------------------

  def configurations(self,levlist=None,imap=False):
    """ Return list of all configurations. """
    if levlist is None: levlist=self.indices()
    if imap is True: levlist=map(self._mapme,levlist)
    out=[]
    for i in levlist:
      if self.v_jconf[i] in out: continue
      out.append(self.v_jconf[i])
    return out

# ------------------------------------------------------------------------------

  def configuration(self,idx,imap=False):
    """ Return configuration of level with given index. """
    tidx=self._mapme(idx,imap)
    if tidx is None: return None
    return self.v_jconf[tidx]

# ------------------------------------------------------------------------------

  def seniority(self,idx,imap=False):
    tidx=self._mapme(idx,imap)
    if tidx is None: return None
    return self.v_seniority[tidx]

# ------------------------------------------------------------------------------

  def j2range(self,levlist=None,imap=False):
    """ Return tuple containing the minimum and maximum 2J values. """
    if levlist is None: levlist=self.indices()
    if imap is True: levlist=map(self.mapme,levlist)
    j2s=[self.v_j2[i] for i in levlist if i is not None]
    return (min(j2s),max(j2s))

# ------------------------------------------------------------------------------

  def j2(self,idx,imap=False):
    tidx=self._mapme(idx,imap)
    if tidx is None: return None
    return self.v_j2[tidx]

# ------------------------------------------------------------------------------

  def jrange(self,levlist=None,imap=False):
    """ Return tuple containing the minimum and maximum 2J values. """
    (j2min,j2max)=self.j2range(levlist,imap)
    return (0.5*float(j2min),0.5*float(j2max))

# ------------------------------------------------------------------------------

  def j(self,idx,imap=False):
    tidx=self._mapme(idx,imap)
    if tidx is None: return None
    return 0.5*float(self.v_j2[tidx])

# ------------------------------------------------------------------------------

  def level(self,idx,imap=False):
    """ Return all information for single level (conf,j2,alpha,par,E) """
    tidx=self._mapme(idx,imap)
    if tidx is None: return (None,None,None,None,None)
    out=(self.v_jconf[tidx],self.v_j2[tidx],self.v_seniority[tidx],\
         self.v_parity[tidx], self.v_energy[tidx])
    return out

# ------------------------------------------------------------------------------

  def label(self,idx,imap=False,spacer=None):
    """ Return level label; optional spacer can be used instead of space. """
    tidx=self._mapme(idx,imap)
    if tidx is None: return None
    j2=self.j2(idx,imap=imap)
    if j2%2 == 0:
      jstr=str(j2/2)
    else:
      jstr=str(j2)+'/2'
    out=str(self.v_jconf[tidx])+' J=%s'%jstr
    alfa=self.v_seniority[tidx]
    if alfa > 1: out+='#%s'%alfa
    if spacer is not None: out=out.replace(' ',spacer)
    return out

# ------------------------------------------------------------------------------

  def weight(self,idx,imap=False):
    """ Return weight of level. """
    tidx=self._mapme(idx,imap)
    if tidx is None: return None
    return self.j2(idx,imap=imap)+1

# ------------------------------------------------------------------------------

  def add_level(self,conf,j2,energy,idx=False):
    """ Add new level to list """
    if idx is not False and idx in self:
      msg='index already exists: %s'%idx
      raise ErrorRAP(msg)
    tidx=idx
    if tidx is False: 
      tidx=1
      if len(self) > 0: tidx=max(self.indices())+1
    tconf=conf
    if not isinstance(tconf,rap_jconfig): tconf=rap_jconfig(tconf)
    if self.num_electrons() is None:
      self.v_nelec=tconf.num_electrons()
    else:
      if tconf.num_electrons() != self.num_electrons():
        msg='inconsistent number of electrons; expect/have: %s, %s'%\
            (self.num_electrons(),tconf.num_electrons())
        raise ErrorRAP(msg)
    self.v_jconf[tidx]=tconf
    self.v_j2[tidx]=j2
    self.v_seniority[tidx]=-1
    self.v_parity[tidx]=tconf.parity()
    self.v_energy[tidx]=energy
    self._reprocess()

# ------------------------------------------------------------------------------

  def add_levels(self,toadd,idxstart=None):
    """ 
    Add set of levels to structure. 

    The input list, toadd, contains tuples with: energy (Ry), index,
    configuration, and 2J.  If the indexing is to be automated, then set index 
    in the tuples to -1.
    """
    if len(toadd) == 0: return
    toadd.sort()
    idxcur=1
    if len(self) > 0: idxcur=max(self.indices())+1
    if idxstart is not None:
      if idxstart < idxcur: 
        msg='invalid starting index: %s'%idxstart
        raise ErrorRAP(msg)
      idxcur=idxstart
    tmp={}
    nechk=None
    for (en,idx,conf,j2) in toadd:
      tidx=idx
      if tidx in self or tidx in tmp: 
        msg='duplicate index: %s'%tidx
        raise ErrorRAP(msg)
      if tidx < 0:
        tidx=idxcur
        idxcur+=1
      tconf=conf
      if not isinstance(conf,rap_jconfig): tconf=rap_jconfig(conf)
      if nechk is None: nechk=tconf.num_electrons()
      if tconf.num_electrons() != nechk:
        msg='inconsistent number of electrons; expected/have: %s, %s'%\
            (nechk,tconf.num_electrons())
        raise ErrorRAP(msg)
      tmp[tidx]=(en,tconf,j2)
    self.v_nelec=nechk
    for idx,(en,conf,j2) in tmp.iteritems():
      self.v_jconf[idx]=conf
      self.v_j2[idx]=j2
      self.v_seniority[idx]=-1
      self.v_parity[idx]=conf.parity()
      self.v_energy[idx]=en
    self._reprocess()

# ------------------------------------------------------------------------------

  def set_lmap(self,other,**args):
    """ Map to other rap_states_jj instance. """
    if self.num_electrons() != other.num_electrons():
      msg='inconsistent number of electrons; expected/have: %s, %s'%\
          (self.num_electrons(),other.num_electrons())
      raise ErrorRAP(msg)
    if isinstance(other,rap_states_jj):
      self._map_jj(other,**args)
    elif isinstance(other,rap_states_lsj):
      self._map_lsj(other,**args)
    elif isinstance(other,rap_states_jp):
      self._map_jp(other,**args)
    else:
      msg='argument must be of type rap_states_jj or rap_states_lsj'
      raise ErrorRAP(msg)

# ------------------------------------------------------------------------------

  def cfilter(self,clist,subset=None,imap=False):
    """ Filter list of states based on configurations in clist. """
    if not isinstance(clist,list) and not isinstance(clist,tuple): clist=[clist]
    tclist=[]
    for c in clist:
      tc=c
      if not isinstance(tc,rap_jconfig): tc=rap_jconfig(tc)
      tclist.append(tc)

    idx=self.indices()
    if subset is not None:
      idx=subset[:]
      if imap is True: idx=self.lmap(idx)

    out=[]
    for i in idx:
      if self.configuration(i) in tclist: out.append(i)
    if imap is True: return self.lmap_reverse(out)
    return out

# ------------------------------------------------------------------------------

  def j2filter(self,j2list,subset=None,imap=False):
    """ Filter list of states based on 2J values in j2list. """
    if not isinstance(j2list,list) and not isinstance(j2list,tuple): 
      j2list=[j2list]
    tj2list=map(int,j2list)

    idx=self.indices()
    if subset is not None:
      idx=subset[:]
      if imap is True: idx=self.lmap(idx)

    out=[]
    for i in idx:
      if self.j2(i) in tj2list: out.append(i)
    if imap is True: return self.lmap_reverse(out)
    return out

# ------------------------------------------------------------------------------

  def jfilter(self,jlist,subset=None,imap=False):
    """ Filter list of states based on J values in j2list. """
    if not isinstance(jlist,list) and not isinstance(jlist,tuple): jlist=[jlist]
    tjlist=map(float,jlist)

    idx=self.indices()
    if subset is not None:
      idx=subset[:]
      if imap is True: idx=self.lmap(idx)

    out=[]
    for i in idx:
      if self.j(i) in tjlist: out.append(i)
    if imap is True: return self.lmap_reverse(out)
    return out

# ------------------------------------------------------------------------------

  def symfilter(self,j2,p=None,subset=None,imap=False):
    """ Filter list of states based on symmetry. """
    if p is None:
      try:
        (jj2,pp)=j2
      except:
        msg='do not understand symmetry'
        raise ErrorRAP(msg)
    else:
      jj2=j2
      pp=p

    idx=self.indices()
    if subset is not None:
      idx=subset[:]
      if imap is True: idx=self.lmap(idx)

    out=[]
    for i in idx:
      tj2=self.j2(i)
      tp=self.parity(i)
      if tj2 != jj2 or tp != pp: continue
      out.append(i)
    if imap is True: return self.lmap_reverse(out)
    return out

# ------------------------------------------------------------------------------

  def search(self,conf=None,j2=None,parity=None,alpha=None,energy=None):
    """ 
    Look for matching levels in structure based on configuration and
    2J values.  If duplicate terms are present and alpha is not given, a list
    of all matching indices are returned.  If alpha is given, the index matching
    that seniority is returned.  If energy is given (and alpha is not), then
    the senority closest to that energy is returned.  If conf and term are 
    both None, then return matching index based on j2, parity, and energy order
    given by alpha.
    """
    if j2 is None: 
      if conf is None:
        msg="if no 2J is specified, must provide configuration"
        raise ErrorRAP(msg)
      levs=self.cfilter(conf)
      return levs

    levs=self.j2filter(j2)
    if conf is None or term is None:
      if conf is not None: parity=conf.parity()
      if parity is None: 
        raise RuntimeError('need parity when conf is None')
      levs=self.pfilter(parity,subset=levs)
      if levs == []: raise RuntimeError('matching level not found')
      if alpha is not None:
        return levs[alpha-1]
      elif energy is not None:
        tmp=[]
        for i in levs: tmp.append((abs(energy-self.energy(i)),i))
        return min(tmp)[1]
      return levs

    levs=self.cfilter(conf,subset=levs)
    if levs == []: raise RuntimeError('matching level not found')
    if len(levs) == 1: return levs[0]
    if alpha is not None:
      for i in levs: 
        if self.seniority(i) == alpha: return i
      raise RuntimeError('matching level not found')
    elif energy is not None:
      tmp=[]
      for i in levs: tmp.append((abs(energy-self.energy(i)),i))
      return min(tmp)[1]
    return levs

# ------------------------------------------------------------------------------

  def load_fac(self,filename,nelec=None):
    """ 
    Load in level information from a FAC energy level file. 

    The optional argument, nelec, will tell the function to look for levels with
    the given number of electrons.  Otherwise, the first level set is taken.
    """
    fin=open(filename,'r')

    # basic info
    z=None
    ne=None
    core=None
    nlev=None
    while 1:
      line=fin.readline()
      if line.strip() == '': break
      if ' Z = ' in line:
        z=int(line.split()[-1])
      elif 'E0' in line:
        e0=float(line.split()[-1])

    # loop over blocks
    done=False
    toadd=[]
    while 1:
      if done: break

      # level set info
      while 1:
        line=fin.readline()
        if line == '':
          done=True
          break
        if 'VNL' in line: break
        if 'NELE' in line:
          ne=int(line.split()[-1])
        elif 'NLEV' in line:
          nlev=int(line.split()[-1])
      if nelec is not None and ne != nelec: continue

      # read in level data
      while 1:
        line=fin.readline()
        if line.strip() == '': break
        tmp=line.split()
        idx=int(tmp[0])
        en=float(tmp[2])/RYDBERG
        j2=int(tmp[5])

        # sub-shell occupancy
        nocc={}
        tmp=re.findall(r'(\d+)\*(\d+)',line)
        for t in tmp: nocc[int(t[0])]=int(t[1])
        ssocc={}
        tmp=re.findall(r'(\d+[a-z])(\d+)',line)
        for t in tmp: ssocc[t[0]]=int(t[1])
        for n,wn in nocc.iteritems():
          wdiff=wn
          noell=True
          for l in range(n):
            ss=rlf_ssencode(n,l)
            if ss not in ssocc: continue
            noell=False
            wdiff-=ssocc[ss]
          if noell:              # one or all sub-shells filled
            if 2*n**2 == wn:
              for l in range(n):
                ss=rlf_ssencode(n,l)
                ssocc[ss]=4*l+2
              continue
            for l in range(n):
              ss=rlf_ssencode(n,l)
              tw=rlf_ssoccupancy(ss)
              done=False
              if not done and tw == wn:   # fill first available
                ssocc[ss]=wn
              else:
                ssocc[ss]=0
            continue
          if wdiff > 0:
            for l in range(n):
              ss=rlf_ssencode(n,l)
              if ss in ssocc: continue
              if rlf_ssoccupancy(ss) <= wdiff:
                ssocc[ss]=rlf_ssoccupancy(ss)
                wdiff-=ssocc[ss]
          if wdiff != 0:                    # sanity check
            msg='inconsistancy in FAC level file?  idx:%i'%idx
            raise ErrorRAP(msg)

        tmp=re.findall(r'([\w+-]+\(\d+\)\d+)',line)
        tconf=' '.join(tmp)
        conf=rap_jconfig(tconf,ne=ne,ssocc=ssocc)

        toadd.append((en,idx,conf,j2))
      if len(toadd) > 0: break
    fin.close()

    if len(toadd) == 0:
      if nelec is None:
        msg="level data found found"
      else:
        pl=''
        if nelec > 1: pl='s'
        msg="level data with %i electron%s not found"%(nelec,pl)
      raise ErrorRAP(msg) 

    self.set_nuclear_charge(z)
    self.set_core_energy(core)
    self.add_levels(toadd)

